package tree

import (
	"fmt"
	"gitee.com/kklt1996/data-structure/queue"
	"gitee.com/kklt1996/data-structure/stack"
	"gitee.com/kklt1996/data-structure/util"
	"strings"
)

/*
	二分搜索树是二叉树,因为其中序列遍历的有序性也称二分搜索树为排序树

	二分搜索树的每隔节点的值:
		大于其左子节点的所有节点的值
		小于其右子节点的所有节点的值
	为了满足此特性,二分搜索树中不允许有相等的元素

	每一棵子树也是二分搜索树

	存入二分搜索树中的节点必须是可以比较的,想要存入基础类型就必须给基础类型实现
	 common.CompareAble接口

	由于二分搜索树可能局部退化成链表,所以改进为带自平衡机制的AVL树
*/
type avlNode struct {
	// 当前节点的值
	key, value interface{}
	// 当前节点的高度
	height int
	// 左右子树
	left, right *avlNode
}

/*
	自平衡二分搜索树
*/
type AVLTree struct {
	// 二分搜索树的根节点
	root *avlNode
	// 元素元素个数
	size int

	/*
		使用comparator比较堆中元素大小,实现common.CompareAble接口和传入比较器形式二选一
		1 thisKey＞compareKey
		0 thisKey=compareKey
		-1 thisKey<compareKey
	*/
	comparator func(thisKey interface{}, compareKey interface{}) int
}

/*
	O(1)
	获取元素个数
*/
func (tree AVLTree) GetSize() int {
	return tree.size
}

/*
	O(1)
	判断是否为空
*/
func (tree AVLTree) IsEmpty() bool {
	return tree.size == 0
}

/*
	O(h)
	添加元素
*/
func (tree *AVLTree) Add(key interface{}, value interface{}) {
	tree.root = tree.add(tree.root, key, value)
}

/*
	向以node为根节点的二分搜索树中插入元素e
	返回以node为根节点的二分搜索树插入新节点之后以的根
	空节点也是一个二叉树
*/
func (tree *AVLTree) add(root *avlNode, key interface{}, value interface{}) *avlNode {
	// 找到满足条件,且为空的位置,就插入元素返回根节点,也就是根节点赋值为新增节点的地址
	if root == nil {
		tree.size++
		// 这就是递归的临界点,在最小的树中找到了插入元素的位置,叶子节点的默认高度为1
		return &avlNode{key: key, value: value, height: 1}
	} else {
		// 要添加的元素和当前根节点比较大小
		i := tree.compare(key, root.key)
		if i < 0 {
			root.left = tree.add(root.left, key, value)
		} else if i > 0 {
			root.right = tree.add(root.right, key, value)
		} else if i == 0 {
			// 重复添加则覆盖节点的value
			root.value = value
		}
		return tree.makeBalance(root)
	}
}

/*
	O(n)
	检查当前节点是不是二分搜索树
*/
func (tree AVLTree) IsBSTTree() bool {
	return tree.isBSTTree(tree.root)
}

func (tree AVLTree) isBSTTree(root *avlNode) bool {
	var preKey interface{}
	isBSTTree := true
	tree.InOrder(func(key interface{}, value interface{}) bool {
		if preKey != nil {
			if tree.compare(key, preKey) < 0 {
				isBSTTree = false
				return true
			}
		}
		preKey = key
		return false
	})
	return isBSTTree
}

/*
	O(n)
	检查一颗树是不是平衡二叉树
*/
func (tree AVLTree) IsBalanceTree() bool {
	return tree.isBalanceTree(tree.root)
}

func (tree AVLTree) isBalanceTree(root *avlNode) bool {
	// 如果是空节点,肯定是平衡的
	if root == nil {
		return true
	}
	// 检查当前节点是否是平衡二叉树
	balanceFactor := tree.getBalanceFactor(root)
	if balanceFactor > 1 || balanceFactor < -1 {
		return false
	} else {
		// 当前节点是平衡二叉树,还要检查左右子树是否是平衡二叉树
		leftBalanceTree := tree.isBalanceTree(root.left)
		rightBalanceTree := tree.isBalanceTree(root.right)
		// 左右子树也是二叉树,根节点才是二叉树
		return leftBalanceTree && rightBalanceTree
	}
}

/*
	O(h)
	查询二分搜索树中是否包含某个元素
*/
func (tree AVLTree) Contains(key interface{}) bool {
	return tree.contains(tree.root, key)
}

/*
	以node为根的二分搜索树是否包含元素value
*/
func (tree AVLTree) contains(root *avlNode, key interface{}) bool {
	// 当查找到节点为空的时候,表明叶子节点都不等于value.那一定是不包含了
	if root == nil {
		return false
	}
	i := tree.compare(key, root.key)
	if i == 0 {
		// 找到表明包含
		return true
	} else if i < 0 {
		// 小于当前节点就去左孩子去继续找
		return tree.contains(root.left, key)
	} else {
		// 大于当前节点就去右边孩子去继续找
		return tree.contains(root.right, key)
	}
}

/*
	查询key对应的value
*/
func (tree AVLTree) Get(key interface{}) interface{} {
	return tree.get(tree.root, key)
}

func (tree AVLTree) get(root *avlNode, key interface{}) interface{} {
	if root == nil {
		return nil
	}
	i := tree.compare(key, root.key)
	if i == 0 {
		return root.value
	} else if i < 0 {
		return tree.get(root.left, key)
	} else {
		return tree.get(root.right, key)
	}
}

/*
	前序遍历,进行一些操作
*/
func (tree AVLTree) PreOrder(operatorFunc func(key interface{}, value interface{}) bool) {
	tree.preOrder(tree.root, operatorFunc)
}

/*
	以某个节点为根节点进行前序列遍历
*/
func (tree AVLTree) preOrder(root *avlNode, operatorFunc func(key interface{}, value interface{}) bool) {
	// 当没有node的时候表示结束
	if root == nil {
		return
	}
	if operatorFunc(root.key, root.value) {
		return
	}
	tree.preOrder(root.left, operatorFunc)
	tree.preOrder(root.right, operatorFunc)
}

/*
	二分搜索树的中序遍历,中序遍历的结果是升序排序的
*/
func (tree AVLTree) InOrder(operatorFunc func(key interface{}, value interface{}) bool) {
	tree.inOrder(tree.root, operatorFunc)
}

/*
	以某个node为根节点进行中序列遍历
*/
func (tree AVLTree) inOrder(root *avlNode, operatorFunc func(key interface{}, value interface{}) bool) {
	// 当没有node的时候表示结束
	if root == nil {
		return
	}
	tree.inOrder(root.left, operatorFunc)
	if operatorFunc(root.key, root.value) {
		return
	}
	tree.inOrder(root.right, operatorFunc)
}

/*
	二分搜索树的后序遍历
*/
func (tree AVLTree) PostOrder(operatorFunc func(key interface{}, value interface{}) bool) {
	tree.postOrder(tree.root, operatorFunc)
}

/*
	以某个node为根节点进行后序列遍历
*/
func (tree AVLTree) postOrder(root *avlNode, operatorFunc func(key interface{}, value interface{}) bool) {
	// 当没有node的时候表示结束
	if root == nil {
		return
	}
	tree.postOrder(root.left, operatorFunc)
	tree.postOrder(root.right, operatorFunc)
	if operatorFunc(root.key, root.value) {
		return
	}
}

/*
	二分搜索树的前序遍历非递归实现
*/
func (tree AVLTree) PreOrderNR(operatorFunc func(key interface{}, value interface{}) bool) {
	var stk stack.Stack = stack.CreateLinkedListStack()
	if tree.root != nil {
		stk.Push(tree.root)
	}

	for !stk.IsEmpty() {
		pop, _ := stk.Pop()
		node := pop.(*avlNode)
		if !operatorFunc(node.key, node.value) {
			return
		}
		if node.right != nil {
			stk.Push(node.right)
		}
		if node.left != nil {
			stk.Push(node.left)
		}
	}

}

/*
	借助队列实现二分搜索树的层序遍历
*/
func (tree AVLTree) LevelOrder(operatorFunc func(key interface{}, value interface{}) bool) {
	var q queue.Queue = queue.CreateLinkedListQueue()
	if tree.root != nil {
		q.Enqueue(tree.root)
	}
	for !q.IsEmpty() {
		// 遍历当前层节点
		dequeue, _ := q.Dequeue()
		node := dequeue.(*avlNode)
		// 搜索结束,直接返回结果
		if operatorFunc(node.key, node.value) {
			return
		}

		// 下一层节点入队,当前层节点遍历完成后就会遍历下一层节点
		if node.left != nil {
			q.Enqueue(node.left)
		}
		if node.right != nil {
			q.Enqueue(node.right)
		}
	}
}

/*
	O(h)
	查找二分搜索树中最小的元素
*/
func (tree AVLTree) Minimum() (interface{}, interface{}, error) {
	if tree.IsEmpty() {
		return nil, nil, AVLIsEmptyError{}
	}
	minimumNode := tree.minimum(tree.root)
	return minimumNode.key, minimumNode.value, nil
}

/*
	递归查找最小元素的节点,也就是微分搜索树最左边的节点
*/
func (tree AVLTree) minimum(node *avlNode) *avlNode {
	if node.left == nil {
		return node
	}
	return tree.minimum(node.left)
}

/*
	O(h)
	查找最小的元素
*/
func (tree AVLTree) Maximum() (interface{}, interface{}, error) {
	if tree.IsEmpty() {
		return nil, nil, AVLIsEmptyError{}
	}
	maximumNode := tree.maximum(tree.root)
	return maximumNode.key, maximumNode.value, nil
}

/*
	递归查找最大元素的节点,也就是微分搜索树最右边的节点
*/
func (tree AVLTree) maximum(node *avlNode) *avlNode {
	if node.right == nil {
		return node
	}
	return tree.maximum(node.right)
}

/*
	O(h)
	删除二分搜索树的最小的节点
*/
func (tree *AVLTree) RemoveMinimum() (interface{}, interface{}, error) {
	if tree.IsEmpty() {
		return nil, nil, AVLIsEmptyError{}
	}
	newRoot, minimumKey, minimumValue := tree.removeMinimum(tree.root)
	// 更新以 tree.root 为根节点子树的根为新的根
	tree.root = newRoot
	return minimumKey, minimumValue, nil
}

/*
	删除以node为根节点的二分搜索树的最小值
	返回删除最小的元素节点后二分搜索树的根和最小的元素
*/
func (tree *AVLTree) removeMinimum(root *avlNode) (*avlNode, interface{}, interface{}) {
	if root.left == nil {
		// 找到最小的元素,获取要返回的右子树
		rightTree := root.right
		// 删除节点和右孩子的关系
		root.right = nil
		// 元素个数-1
		tree.size--
		// 返回最小节点的右子树,和最小节点的值
		return rightTree, root.key, root.value
	} else {
		// 删除左子树最小的元素
		newRoot, minimumKey, minimumValue := tree.removeMinimum(root.left)
		// 更新以 root.left 为根节点子树的根为新的根
		root.left = newRoot
		// 返回原根和删除最小节点的value
		return tree.makeBalance(root), minimumKey, minimumValue
	}
}

/*
	O(h)
	删除二分搜索树的最大的节点
*/
func (tree *AVLTree) RemoveMaximum() (interface{}, interface{}, error) {
	if tree.IsEmpty() {
		return nil, nil, AVLIsEmptyError{}
	}
	newRoot, maximumKey, maximumValue := tree.removeMaximum(tree.root)
	tree.root = newRoot
	return maximumKey, maximumValue, nil
}

/*
	删除以node为根节点的二分搜索树的最大值的节点
	返回删除最大的元素节点后二分搜索树的根
*/
func (tree *AVLTree) removeMaximum(root *avlNode) (*avlNode, interface{}, interface{}) {
	if root.right == nil {
		// 找到最大的元素,获取要返回的左子树
		leftTree := root.left
		// 删除节点和左孩子的关系
		root.left = nil
		// 元素个数-1
		tree.size--
		// 返回最大节点的左子树
		return leftTree, root.key, root.value
	} else {
		// 删除右子树最大的元素
		newRoot, maximumKey, maximumValue := tree.removeMaximum(root.right)
		root.right = newRoot
		return tree.makeBalance(root), maximumKey, maximumValue
	}
}

/*
	O(h)
	删除二分搜索树中任意节点
*/
func (tree *AVLTree) RemoveElement(key interface{}) (interface{}, error) {
	if tree.IsEmpty() {
		return nil, AVLIsEmptyError{}
	}
	var retValue interface{}
	tree.root, retValue = tree.removeElement(tree.root, key)
	return retValue, nil
}

/*
	删除以某个node为根节点的二分搜索树中的某一个节点
	返回删除指定节点后新二分搜索树的根
*/
func (tree *AVLTree) removeElement(root *avlNode, key interface{}) (*avlNode, interface{}) {
	// 根节点为空直接返回
	if root == nil {
		return nil, nil
	}
	i := tree.compare(key, root.key)
	// 当前子树删除元素后新的根节点
	var retNode *avlNode
	var retValue interface{}
	if i == 0 {
		// 递归终止条件
		// 找到需要删除的节点
		if root.left == nil {
			// 左子树为空,那么直接将右节点升级为二分搜索树的根
			rightTree := root.right
			root.right = nil
			tree.size--
			return rightTree, root.value
		} else if root.right == nil {
			// 右子树为空,那么直接将左子树升级为二分搜索树的根
			leftTree := root.left
			root.left = nil
			tree.size--
			return leftTree, root.value
		} else {
			// 左右子树都不为空,那么为满足二分搜索树的性质
			// 需要从右子树中删除最小的节点或者从左子树中删除最大的节点
			// 将删除的节点赋值给当前根节点
			rightNewRoot, removeKey, removeValue := tree.removeMinimum(root.right)
			root.right = rightNewRoot
			retNode = root
			retValue = root.value
			root.key = removeKey
			root.value = removeValue
		}
	} else {
		// 用更小的结果拼凑最终的结果,在子树中删除元素
		if i < 0 {
			root.left, retValue = tree.removeElement(root.left, key)
		} else {
			root.right, retValue = tree.removeElement(root.right, key)
		}
		// 根不变
		retNode = root
	}
	// 返回自平衡后的根
	return tree.makeBalance(retNode), retValue
}

/*
	对左右子树发生变化的二叉树进行自平衡操作
	只有左右子树发生变化的,才肯能会引起高度和平衡因子的变化
*/
func (tree *AVLTree) makeBalance(retNode *avlNode) *avlNode {
	// 对删除节点后的子树进行自平衡

	// 维护树的高度
	tree.updateRootHeight(retNode)
	// 计算节点的平衡因子,维护二分搜索树的平衡性
	balanceFactor := tree.getBalanceFactor(retNode)

	// 节点不平衡
	if balanceFactor > 1 || balanceFactor < -1 {
		if balanceFactor > 1 && tree.getBalanceFactor(retNode.left) >= 0 {
			// LL
			// 左子树高度比右子树高度大2.左子树的左子树比右子树高度大1.进行右旋转操作
			// 返回右旋转后新的根
			return tree.rightRotate(retNode)
		} else if balanceFactor < -1 && tree.getBalanceFactor(retNode.right) <= 0 {
			// RR
			// 右子树高度比左子树高度大2.右子树的右子树比左子树高度大1,进行左旋转操作
			return tree.leftRotate(retNode)
		} else if balanceFactor > 1 && tree.getBalanceFactor(retNode.left) < 0 {
			// LR
			// 先对左节点进行左旋,使得整体转换为LL形式
			retNode.left = tree.leftRotate(retNode.left)
			// 对LL形式进行处理
			return tree.rightRotate(retNode)
		} else if balanceFactor < -1 && tree.getBalanceFactor(retNode.right) > 0 {
			// RL
			// 先对右节点进行右旋,使得整体转转为RR形式
			retNode.right = tree.rightRotate(retNode.right)
			return tree.leftRotate(retNode)
		}
	}
	return retNode
}

/*
	获取当前节点的高度
*/
func (tree AVLTree) getHeight(node *avlNode) int {
	if node == nil {
		return 0
	}
	return node.height
}

/*
	根据左右子节点的高度更新当前节点高度
*/
func (tree AVLTree) updateRootHeight(root *avlNode) {
	if root == nil {
		return
	}
	leftTreeHeight := tree.getHeight(root.left)
	rightTreeHeight := tree.getHeight(root.right)
	if leftTreeHeight > rightTreeHeight {
		root.height = leftTreeHeight + 1
	} else {
		root.height = rightTreeHeight + 1
	}
}

/*
	计算节点的平衡因子
*/
func (tree AVLTree) getBalanceFactor(node *avlNode) int {
	if node == nil {
		return 0
	}
	return tree.getHeight(node.left) - tree.getHeight(node.right)
}

/*
		右旋转操作
		 y				 x
		/ \				/ \
	   x   T4------>   z   y
      / \             / \  / \
	 z   T3          T1 T2 T3 T4
    / \
   T1 T2
	x,y,z 节点都可能有左右子树
*/
func (tree AVLTree) rightRotate(root *avlNode) *avlNode {
	// 获取要返回的根节点
	newRoot := root.left
	root.left.right, root.left = root, root.left.right
	// 回溯更新左右孩子节点发生变化的节点
	tree.updateRootHeight(newRoot.right)
	tree.updateRootHeight(newRoot)
	return newRoot
}

/*
	左旋转操作
			y					x
			 \				   / \
			  x     ----->    y   z
               \
                z
*/
func (tree AVLTree) leftRotate(root *avlNode) *avlNode {
	newRoot := root.right
	root.right.left, root.right = root, root.right.left
	tree.updateRootHeight(newRoot.left)
	tree.updateRootHeight(newRoot)
	return newRoot
}

func (tree AVLTree) compare(thisValue interface{}, compareValue interface{}) int {
	if tree.comparator != nil {
		return tree.comparator(thisValue, compareValue)
	} else {
		return util.DefaultComparator(thisValue, compareValue)
	}
}

func (tree AVLTree) String() string {
	res := "["
	tree.inOrder(tree.root, func(key interface{}, value interface{}) bool {
		res += fmt.Sprint(key) + ","
		return false
	})
	res = strings.TrimSuffix(res, ",")
	res += "]"
	return res
}

type AVLIsEmptyError struct {
}

func (AVLIsEmptyError) Error() string {
	return "avl is empty"
}
